{ "cells": [ { "cell_type": "markdown", "id": "b2bb0357", "metadata": {}, "source": [ "### PEAQ VAE" ] }, { "cell_type": "code", "execution_count": null, "id": "e0db834c", "metadata": {}, "outputs": [], "source": [ "import os\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n", "os.environ[\"PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION\"] = \"python\"" ] }, { "cell_type": "code", "execution_count": null, "id": "9afebefb", "metadata": {}, "outputs": [], "source": [ "# !cat /home/minz/glockenspiel/descript-audio-codec/conf/v4/25hz_vae_peaq.yml" ] }, { "cell_type": "code", "execution_count": null, "id": "c97a26a7", "metadata": {}, "outputs": [], "source": [ "from suno_utils.audio import Audio\n", "from suno_utils.tasks.dac_vae_peaq import (\n", " preload_models as preload_codec_models,\n", " encode as codec_encode,\n", " decode as codec_decode,\n", ")\n", "\n", "_ = preload_codec_models(\"/app/suno/checkpoints/dac_mw/mw_vae_peaq_128_fix/best/dac/weights.pth\")\n", "# _ = preload_codec_models(\"/app/suno/checkpoints/dac_mw/mw_vae_peaq/best/dac/weights.pth\")\n", "# _ = preload_codec_models(\"s3://suno-data/georg/models/codec/peaq_vae.pth\")" ] }, { "cell_type": "code", "execution_count": null, "id": "e46745a7", "metadata": {}, "outputs": [], "source": [ "audio = Audio.from_file(\"samples/halo_orig.mp3\", n_channels=2, sample_rate=48_000, byte_width=2)\n", "# audio = Audio.from_file(\"samples/first_cig_orig.mp3\", n_channels=2, sample_rate=48_000, byte_width=2)\n", "# audio = Audio.from_file(\"samples/orchestra_orig.mp3\", n_channels=2, sample_rate=48_000, byte_width=2)\n", "# audio = Audio.from_file(\"samples/left_right_orig.mp3\", n_channels=2, sample_rate=48_000, byte_width=2)" ] }, { "cell_type": "code", "execution_count": null, "id": "0963532b", "metadata": {}, "outputs": [], "source": [ "out = codec_encode(audio)" ] }, { "cell_type": "code", "execution_count": null, "id": "4363c4d5", "metadata": {}, "outputs": [], "source": [ "codec_decode(out).play()" ] }, { "cell_type": "code", "execution_count": null, "id": "a54f4208", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "95abcbab", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "1d9f2349", "metadata": {}, "source": [ "### Oobleck VAE" ] }, { "cell_type": "code", "execution_count": null, "id": "9d1d0ed9", "metadata": {}, "outputs": [], "source": [ "import json\n", "from stable_audio_tools.models.autoencoders import create_autoencoder_from_config\n", "# from suno_utils.models.stable_audio_models.autoencoders import create_autoencoder_from_config\n", "\n", "BASE_DIR = \"/home/christian/christian/stable-audio-tools\"" ] }, { "cell_type": "code", "execution_count": null, "id": "3b9580f7", "metadata": {}, "outputs": [], "source": [ "config_fp = \"stable_audio_tools/configs/model_configs/txt2audio/stable_audio_2_0_mert_48khz.json\"\n", "with open(os.path.join(BASE_DIR, config_fp)) as f:\n", " config = json.load(f)\n", "assert(config[\"sample_rate\"] == 48_000)\n", "ae_config = {\n", " \"model\": config[\"model\"][\"pretransform\"][\"config\"],\n", " \"sample_rate\": 48_000,\n", "}" ] }, { "cell_type": "code", "execution_count": null, "id": "0fe299d5", "metadata": {}, "outputs": [], "source": [ "model = create_autoencoder_from_config(ae_config)" ] }, { "cell_type": "code", "execution_count": null, "id": "76153006", "metadata": {}, "outputs": [], "source": [ "ckpt_fp = os.path.join(BASE_DIR, \"checkpoints/vae/vae_model_unwrap-epoch=52-step=1140000.ckpt\")\n", "sd = torch.load(ckpt_fp, map_location=\"cpu\")\n", "model.load_state_dict(sd[\"state_dict\"]);" ] }, { "cell_type": "code", "execution_count": null, "id": "c2fe3bf9", "metadata": {}, "outputs": [], "source": [ "# audio = Audio.from_file(\"samples/halo_orig.mp3\", n_channels=2, sample_rate=48_000, byte_width=2)\n", "# audio = Audio.from_file(\"samples/first_cig_orig.mp3\", n_channels=2, sample_rate=48_000, byte_width=2)\n", "# audio = Audio.from_file(\"samples/orchestra_orig.mp3\", n_channels=2, sample_rate=48_000, byte_width=2)\n", "audio = Audio.from_file(\"samples/left_right_orig.mp3\", n_channels=2, sample_rate=48_000, byte_width=2)" ] }, { "cell_type": "code", "execution_count": null, "id": "7b804ebf", "metadata": {}, "outputs": [], "source": [ "in_arr = torch.from_numpy(audio.array_float)\n", "with torch.no_grad():\n", " out = model.encode(in_arr[None])\n", "out.shape" ] }, { "cell_type": "code", "execution_count": null, "id": "b39e7c32", "metadata": {}, "outputs": [], "source": [ "with torch.no_grad():\n", " out2 = model.decode(out)\n", "out_audio = Audio.from_array_float(out2.detach().cpu().numpy()[0], sample_rate=48_000)\n", "out_audio.play()" ] }, { "cell_type": "code", "execution_count": null, "id": "6506b06e", "metadata": {}, "outputs": [], "source": [ "# out_audio.to_hq_mp3(\"tmp/oob_halo.mp3\")\n", "# out_audio.to_hq_mp3(\"tmp/oob_first_cig.mp3\")\n", "# out_audio.to_hq_mp3(\"tmp/oob_orchestra.mp3\")\n", "# out_audio.to_hq_mp3(\"tmp/oob_left_right.mp3\")" ] }, { "cell_type": "code", "execution_count": null, "id": "6b330b28", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "d9397402", "metadata": {}, "source": [ "## Custom code oobleck" ] }, { "cell_type": "code", "execution_count": null, "id": "b2b0496f", "metadata": {}, "outputs": [], "source": [ "import os\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n", "os.environ[\"PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION\"] = \"python\"" ] }, { "cell_type": "code", "execution_count": null, "id": "f8e91f06", "metadata": {}, "outputs": [], "source": [ "from suno_utils.audio import Audio\n", "from suno_utils.tasks.oobleck import encode, preload_models, decode" ] }, { "cell_type": "code", "execution_count": null, "id": "a59f58c4", "metadata": {}, "outputs": [], "source": [ "!aws s3 cp /home/christian/christian/stable-audio-tools/stable_audio_tools/configs/model_configs/txt2audio/stable_audio_2_0_mert_48khz.json s3://suno-data/georg/models/codec/oobleck_23hz.ckpt" ] }, { "cell_type": "code", "execution_count": null, "id": "991815b3", "metadata": {}, "outputs": [], "source": [ "BASE_DIR = \"/home/christian/christian/stable-audio-tools\"\n", "_ = preload_models(\n", " os.path.join(BASE_DIR, \"checkpoints/vae/vae_model_unwrap-epoch=52-step=1140000.ckpt\"), \n", " os.path.join(BASE_DIR, \"stable_audio_tools/configs/model_configs/txt2audio/stable_audio_2_0_mert_48khz.json\"),\n", ")" ] }, { "cell_type": "code", "execution_count": null, "id": "05ba57ef", "metadata": {}, "outputs": [], "source": [ "audio = Audio.from_file(\"samples/halo_orig.mp3\", n_channels=2, sample_rate=48_000, byte_width=2)\n", "# audio = Audio.from_file(\"samples/first_cig_orig.mp3\", n_channels=2, sample_rate=48_000, byte_width=2)\n", "# audio = Audio.from_file(\"samples/orchestra_orig.mp3\", n_channels=2, sample_rate=48_000, byte_width=2)\n", "# audio = Audio.from_file(\"samples/left_right_orig.mp3\", n_channels=2, sample_rate=48_000, byte_width=2)" ] }, { "cell_type": "code", "execution_count": null, "id": "3ee11e75", "metadata": {}, "outputs": [], "source": [ "out = encode(audio)" ] }, { "cell_type": "code", "execution_count": null, "id": "eef30cef", "metadata": {}, "outputs": [], "source": [ "decode(out).play()" ] }, { "cell_type": "code", "execution_count": null, "id": "f2e83568", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "248d5f18", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "ad11b225", "metadata": {}, "source": [ "### Modal embed oobleck" ] }, { "cell_type": "code", "execution_count": null, "id": "55a8216e", "metadata": {}, "outputs": [], "source": [ "# modal run ~/code/glockenspiel/suno_utils/suno_utils/scripts/gpt/modal_encode.py \\\n", "# --base-s3-dir='s3://suno-data/datasets/bundles/v1/genius_hq' \\\n", "# --embed-type='oobleck_23' \\\n", "# --chunksize='500' \\\n", "# --min-duration-s='60' \\\n", "# --max-duration-s='480' \\\n", "# --first-only=True \\\n", "# --force-overwrite=True \\\n", "# --output-name='oobleck_23_v0'\n", "\n", "# modal run ~/code/glockenspiel/suno_utils/suno_utils/scripts/gpt/modal_encode.py \\\n", "# --base-s3-dir='s3://suno-data/datasets/bundles/v1/youtube_music' \\\n", "# --embed-type='dac_gvq' \\\n", "# --chunksize='500' \\\n", "# --min-duration-s='60' \\\n", "# --max-duration-s='480' \\\n", "# --output-name='oobleck_23_v0'\n", "\n", "# --first-only=True \\" ] }, { "cell_type": "code", "execution_count": null, "id": "5b404815", "metadata": {}, "outputs": [], "source": [ "# 2.5tb for 100k hours" ] }, { "cell_type": "code", "execution_count": null, "id": "cc020ba6", "metadata": {}, "outputs": [], "source": [ "# !aws s3 cp s3://suno-data/datasets/bundles/v1/genius_hq/oobleck_23_v0/part_0.npz tmp/" ] }, { "cell_type": "code", "execution_count": null, "id": "1f5e3c83", "metadata": {}, "outputs": [], "source": [ "# TODO: output volume issue..." ] }, { "cell_type": "code", "execution_count": null, "id": "ce95a134", "metadata": {}, "outputs": [], "source": [ "# import numpy as np\n", "# aa = np.load(\"tmp/part_0.npz\")" ] }, { "cell_type": "code", "execution_count": null, "id": "62af5b20", "metadata": {}, "outputs": [], "source": [ "# aa[list(aa.keys())[0]].shape" ] }, { "cell_type": "code", "execution_count": null, "id": "a95571bf", "metadata": {}, "outputs": [], "source": [ "# decode(aa[list(aa.keys())[0]][:2000,:]).play()" ] }, { "cell_type": "code", "execution_count": null, "id": "95c4e964", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "1abe2340", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "b81bc2fe", "metadata": {}, "source": [ "### Modal encode dac_vae_peaq" ] }, { "cell_type": "code", "execution_count": null, "id": "d975c269", "metadata": {}, "outputs": [], "source": [ "# modal run ~/code/glockenspiel/suno_utils/suno_utils/scripts/gpt/modal_encode.py \\\n", "# --base-s3-dir='s3://suno-data/datasets/bundles/v1/genius_hq' \\\n", "# --embed-type='dac_vae' \\\n", "# --chunksize='500' \\\n", "# --min-duration-s='60' \\\n", "# --max-duration-s='480' \\\n", "# --force-overwrite=True \\\n", "# --output-name='dac_vae_128'\n", "\n", "# modal run ~/code/glockenspiel/suno_utils/suno_utils/scripts/gpt/modal_encode.py \\\n", "# --base-s3-dir='s3://suno-data/datasets/bundles/v1/youtube_music' \\\n", "# --embed-type='dac_vae' \\\n", "# --chunksize='500' \\\n", "# --min-duration-s='60' \\\n", "# --max-duration-s='480' \\\n", "# --output-name='dac_vae_128'\n", "\n", "# --force-overwrite=True \\\n", "# --first-only=True \\" ] }, { "cell_type": "code", "execution_count": null, "id": "643bb48d", "metadata": {}, "outputs": [], "source": [ "# !aws s3 cp s3://suno-data/datasets/bundles/v1/genius_hq/dac_vae_128/part_0.npz tmp/" ] }, { "cell_type": "code", "execution_count": null, "id": "714e8391", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "aa = np.load(\"tmp/part_0.npz\")" ] }, { "cell_type": "code", "execution_count": null, "id": "34e25f75", "metadata": {}, "outputs": [], "source": [ "from suno_utils.audio import Audio\n", "from suno_utils.tasks.dac_vae_peaq import (\n", " preload_models as preload_codec_models,\n", " encode as codec_encode,\n", " decode as codec_decode,\n", ")\n", "_ = preload_codec_models(\"s3://suno-data/georg/models/codec/dac_vae_128.pth\")" ] }, { "cell_type": "code", "execution_count": null, "id": "e115a536", "metadata": {}, "outputs": [], "source": [ "aa[list(aa.keys())[0]].shape" ] }, { "cell_type": "code", "execution_count": null, "id": "4eb522d5", "metadata": {}, "outputs": [], "source": [ "codec_decode(aa[list(aa.keys())[0]][:,:]).play()" ] }, { "cell_type": "code", "execution_count": null, "id": "0184bf67", "metadata": {}, "outputs": [], "source": [ "codec_decode(aa[list(aa.keys())[1]][:,:]).play()" ] }, { "cell_type": "code", "execution_count": null, "id": "baae25e1", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "0c3500ca", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "d13e5d0f", "metadata": {}, "source": [ "### Test diffusion model" ] }, { "cell_type": "code", "execution_count": null, "id": "010f9715", "metadata": {}, "outputs": [], "source": [ "%matplotlib inline\n", "from matplotlib import pyplot as plt" ] }, { "cell_type": "code", "execution_count": null, "id": "33481a56", "metadata": {}, "outputs": [], "source": [ "import os\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n", "os.environ[\"PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION\"] = \"python\"" ] }, { "cell_type": "code", "execution_count": null, "id": "c0f5a102", "metadata": {}, "outputs": [], "source": [ "import sys\n", "sys.path.insert(0, \"/home/georg/code/christian/stable-audio-tools\")" ] }, { "cell_type": "code", "execution_count": null, "id": "6319b7da", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "import torch\n", "from suno_utils.audio import Audio\n", "from suno_utils.tasks.mert_25 import (\n", " preload_models as preload_semantic_models,\n", " encode as semantic_encode,\n", ")\n", "from suno_utils.tasks.dac_2c_12cb import (\n", " preload_models as preload_codec_models,\n", " encode as codec_encode,\n", ")\n", "from suno_utils.utils.text import read_json\n", "_ = preload_semantic_models(\n", " checkpoint_filepath=\"s3://suno-data/georg/models/semantic/mert_25.pt\",\n", " centroids_filepath=\"s3://suno-data/georg/models/semantic/mert_25_2x4k.npy\",\n", ")\n", "_ = preload_codec_models(\n", " checkpoint_filepath=\"s3://suno-data/georg/models/codec/dac_2c_25x12.pt\",\n", ")\n", "BASE_DIR = \"/home/christian/christian/stable-audio-tools\"" ] }, { "cell_type": "code", "execution_count": null, "id": "e2839ef3", "metadata": {}, "outputs": [], "source": [ "from stable_audio_tools.interface.gradio import load_model\n", "from stable_audio_tools.inference.generation import upsample_diffusion\n", "from stable_audio_tools.models.utils import apply_normalization\n", "from stable_audio_tools.interface.gradio import load_model" ] }, { "cell_type": "code", "execution_count": null, "id": "ca7f8ae1", "metadata": {}, "outputs": [], "source": [ "# !CUDA_VISIBLE_DEVICES=1 python /home/christian/christian/stable-audio-tools/unwrap_model.py \\\n", "# --model-config /home/christian/christian/stable-audio-tools/stable_audio_tools/configs/model_configs/txt2audio/stable_audio_2_0_mert_48khz.json \\\n", "# --ckpt-path /home/christian/christian/stable-audio-tools/harmonai_train/drn8j5c3/checkpoints/epoch=116-step=120000.ckpt \\\n", "# --name mert_condition_large" ] }, { "cell_type": "code", "execution_count": null, "id": "1f020a0a", "metadata": {}, "outputs": [], "source": [ "# !CUDA_VISIBLE_DEVICES=1 python /home/christian/christian/stable-audio-tools/unwrap_model.py \\\n", "# --model-config /home/christian/christian/stable-audio-tools/stable_audio_tools/configs/model_configs/txt2audio/stable_audio_2_0_codec_48khz.json \\\n", "# --ckpt-path /home/christian/christian/stable-audio-tools/harmonai_train/2zp4baao/checkpoints/epoch=38-step=320000.ckpt \\\n", "# --name codec_condition" ] }, { "cell_type": "code", "execution_count": null, "id": "2d4b0ba8", "metadata": {}, "outputs": [], "source": [ "# !cat /home/christian/christian/stable-audio-tools/stable_audio_tools/configs/model_configs/txt2audio/stable_audio_2_0_codec_48khz.json" ] }, { "cell_type": "code", "execution_count": null, "id": "bc131b59", "metadata": {}, "outputs": [], "source": [ "# # 8x semantic conditioned\n", "# model_type = \"mert\"\n", "# ckpt_fp = \"/home/georg/notebooks/diffusion/mert_condition_large.ckpt\"\n", "# config_fp = \"stable_audio_tools/configs/model_configs/txt2audio/stable_audio_2_0_mert_48khz.json\"\n", "\n", "# 1x semantic confitioned\n", "model_type = \"codec\"\n", "ckpt_fp = \"/home/georg/notebooks/diffusion/codec_condition.ckpt\"\n", "config_fp = \"stable_audio_tools/configs/model_configs/txt2audio/stable_audio_2_0_codec_48khz.json\"\n", "\n", "config_fp = os.path.join(BASE_DIR, config_fp)\n", "tmp_model_config = read_json(config_fp)\n", "model, model_config = load_model(\n", " tmp_model_config,\n", " ckpt_fp,\n", " # pretrained_name=pretrained_name,\n", " # pretransform_ckpt_path=vae_fp,\n", " # model_half=model_half,\n", " device=\"cuda\" if torch.cuda.is_available() else \"cpu\",\n", ")\n", "device = next(model.parameters()).device" ] }, { "cell_type": "code", "execution_count": null, "id": "71eb17e1", "metadata": {}, "outputs": [], "source": [ "audio = Audio.from_file(\"samples/halo_orig.mp3\", n_channels=2, sample_rate=48_000, byte_width=2)\n", "# audio = Audio.from_file(\"samples/first_cig_orig.mp3\", n_channels=2, sample_rate=48_000, byte_width=2)\n", "# audio = Audio.from_file(\"samples/orchestra_orig.mp3\", n_channels=2, sample_rate=48_000, byte_width=2)\n", "# audio = Audio.from_file(\"samples/left_right_orig.mp3\", n_channels=2, sample_rate=48_000, byte_width=2)" ] }, { "cell_type": "code", "execution_count": null, "id": "34f9f293", "metadata": {}, "outputs": [], "source": [ "start_frame = 0\n", "\n", "audio_in = torch.from_numpy(audio.array_float)[:, start_frame : start_frame + 524288]\n", "audio_in = apply_normalization(audio_in, 48000, target_loudness_lufs_db=-16.0)\n", "audio_segment = audio.get_segment(from_s=start_frame/48_000, to_s=(start_frame+524288)/48_000)\n", "if model_type == \"mert\":\n", " latents = semantic_encode(audio_segment)[:, :1]\n", "elif model_type == \"codec\":\n", " latents = codec_encode(audio_segment)\n", "else:\n", " raise NotImplementedError()\n", "print(latents.shape)\n", "codes = torch.from_numpy(latents.astype(np.int64)).long().T.to(device)\n", "print(codes.shape)" ] }, { "cell_type": "code", "execution_count": null, "id": "6bd4ed74", "metadata": {}, "outputs": [], "source": [ "upsampled = upsample_diffusion(\n", " model,\n", " audio_in,\n", " model_type,\n", " codes=codes,\n", " steps=500,\n", " cfg_scale=1,\n", " sample_size=audio_in.shape[-1],\n", " sample_rate=48_000,\n", ")\n", "upsampled = upsampled.cpu()\n", "upsampled /= upsampled.abs().max()\n", "print(upsampled.shape)\n", "out_audio = Audio.from_array_float(upsampled.detach().cpu().numpy()[0], 48_000)\n", "out_audio.play()" ] }, { "cell_type": "code", "execution_count": null, "id": "03fdc7bb", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "25dbc575", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "367b01e6", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "44e5ad2b", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "e6acf575", "metadata": {}, "source": [ "## Playground" ] }, { "cell_type": "code", "execution_count": null, "id": "4db7fe01", "metadata": {}, "outputs": [], "source": [ "# semantic_data = np.memmap(\"semantic_data.bin\", dtype=np.uint16, mode=\"w+\", shape=(100, 250))\n", "# semantic_data[:] = 0\n", "# semantic_data.flush()\n", "\n", "# codec_data = np.memmap(\"codec_data.bin\", dtype=np.float32, mode=\"w+\", shape=(100, 250, 128))\n", "# codec_data[:] = 0\n", "# codec_data.flush()" ] }, { "cell_type": "code", "execution_count": null, "id": "b39256ef", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "bb858614", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "4cd525c6", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.12" } }, "nbformat": 4, "nbformat_minor": 5 }